import sys, os
import cv2, shutil
import numpy as np

from glob import glob
from os.path import splitext, basename
from src.utils import im2single
from src.keras_utils import load_model, detect_lp
from src.label import Shape
import time
from src.drawing_utils import draw_losangle


class LpDetector():
    def __init__(self, model_path, lp_threshold=.0):
        self.lp_threshold = lp_threshold
        self.wpod_net = load_model(model_path)

    def detect(self, img_path, Ivehicle, im_save=False, im_save_dir='out', verbose=False):
        # Ivehicle = cv2.imread(img_path)  Image of Vehicle
        im_h = Ivehicle.shape[0]
        im_w = Ivehicle.shape[1]
        ratio = float(max(Ivehicle.shape[:2])) / min(Ivehicle.shape[:2])
        side = int(ratio * 288.)
        bound_dim = min(side + (side % (2 ** 4)), 608)
        start = time.time()
        Llp, LlpImgs, _ = detect_lp(self.wpod_net, im2single(Ivehicle), bound_dim, 2 ** 4, (240, 80), self.lp_threshold,
                                    verbose=verbose)
        if verbose:
            print(" wpod net cost : {}".format(time.time() - start))

        RED = (0, 0, 255)
        if len(LlpImgs):
            Ilp = LlpImgs[0]
            I = Ivehicle
            pts = Shape(Llp[0].pts).pts
            ptspx = pts * np.array(I.shape[1::-1], dtype=float).reshape(2, 1)
            draw_losangle(I, ptspx, RED, 3)
            prob = Llp[0]._Label__prob
            if im_save:
                if os.path.exists(im_save_dir) is False:
                    os.makedirs(im_save_dir)
                cv2.putText(I, "CONFIDENCE:" + str(prob), (0, 70), cv2.FONT_HERSHEY_COMPLEX, 2, (0, 0, 255), 5)
                cv2.imwrite('%s/%s_full.png' % (im_save_dir, os.path.basename(img_path)), I)
                cv2.imwrite('%s/%s_lp.png' % (im_save_dir, os.path.basename(img_path)),
                            np.array(Ilp * 255., dtype=np.uint8))

            # cv format image BGR
            return True, ptspx, np.array(Ilp * 255., dtype=np.uint8), prob
        else:
            print(" license plate not found in:{}".format(img_path))
            return False, None, None, None


def case1():
    # im_path = './samples/test/03016.jpg'
    # im_path = '/home/leo/Downloads/datas/road_shot_lp/川A0CP01_1.jpg'
    im_path = '/home/leo/Downloads/datas/gen_wpod_det_trainset/mix/val_hardcase2/WJ粤0033X_0car.jpg'

    Ivehicle = cv2.imread(im_path)
    lp_det.detect(im_path, Ivehicle, im_save=True)


def case2(dir=''):
    file_list = []
    for file in os.listdir(dir):
        if os.path.splitext(file)[1] in ['.jpg', '.png']:
            file_list.append(os.path.join(dir, file))

    for i, im_path in enumerate(file_list):
        sys.stdout.write('\r>> loading label %d/%d ' % (
            i+1, len(file_list)))
        sys.stdout.flush()
        if os.path.splitext(im_path)[1] in ['.jpg', '.png']:
            Ivehicle = cv2.imread(im_path)
            lp_det.detect(im_path, Ivehicle, im_save=True)


if __name__ == '__main__':
    lp_det = LpDetector('./training_models/hardcase-network_backup.h5')
    # case1()
    # case2(dir='/home/leo/Downloads/datas/yolo_tiny_lp_det_data_set/val')
    case2(dir='/home/leo/Downloads/datas/gen_wpod_det_trainset/mix/val_hardcase2/')
